import logging
import numpy as np

from util.util_log import setup_logging
from FitnessMultiBeamTrace import FitnessMultiBeamTrace


# 配置日志，默认打印到控制台，也可以设置打印到文件
setup_logging()
# setup_logging(log_file="../files/multi-beam-trace/beam2/QE/L=L1+L4/trace_2025-03-19-1644-16beam-qe.txt")
# 获取日志记录器并记录日志
logger = logging.getLogger("[BIVY]")


class Plant:
    def __init__(self, position, cost_function):
        self.position = position  # 位置（二进制），形状为 (rows, cols)
        self.cost_function = cost_function  # 成本函数
        self.cost = self.calculate_cost()  # 计算成本

    def calculate_cost(self):
        return self.cost_function(self.position)  # 计算成本

    def update_position(self, new_position):
        self.position = new_position  # 更新位置
        self.cost = self.calculate_cost()  # 重新计算成本


class BinaryIvyOptimization:
    def __init__(self, N, max_iterations, rows, cols, fobj):
        self.N = N  # 种群大小
        self.max_iterations = max_iterations  # 最大迭代次数
        self.rows = rows  # 行数
        self.cols = cols  # 列数
        self.fobj = fobj  # 成本函数

        self.var_size = (rows, cols)  # 决策变量矩阵大小

        self.pop = []  # 初始种群数组
        self.best_costs = np.zeros(self.max_iterations)  # 最佳成本历史
        self.convergence_curve = []  # 收敛曲线


    def __generate_population_part(self, before_array, current_array, max_num):
        part = []
        # 添加before和current作为前两个元素
        part.append(before_array.copy().astype(int))
        part.append(current_array.copy().astype(int))
        # 计算 diff 数组
        diff = before_array != current_array
        # 获取所有 False 的坐标
        false_coords = np.argwhere(diff)
        # 生成剩余的个体
        for _ in range(max_num - 2):
            # 复制 current_array 并强制转换为 int 类型
            modified_array = np.copy(current_array).astype(int)
            # 随机选择一些差异行
            if np.sum(diff == True) > 0:
                # 随机选择一部分 False 坐标
                num_to_select = np.random.randint(0, len(false_coords) + 1)
                selected_false_coords = false_coords[np.random.choice(len(false_coords), num_to_select, replace=False)]
                # 将 modified_array 中选取坐标替换为 before_array
                for coord in selected_false_coords:
                    row, col = coord
                    modified_array[row, col] = before_array[row, col]
                # print("Modified array:", modified_array.tolist())
            part.append(modified_array)
        return part

    def initialize_population(self):
        for _ in range(self.N):
            position = np.random.randint(2, size=self.var_size).astype(int)  # 初始化二进制位置
            plant = Plant(position, self.fobj)  # 创建植物对象
            self.pop.append(plant)

    def initialize_population_group(self, phaseBit_mix_ps_before, phaseBit_mix_ps, phaseBit_mix_qe_before,
                                    phaseBit_mix_qe):
        parts_ps = self.__generate_population_part(phaseBit_mix_ps_before, phaseBit_mix_ps, self.N // 2)
        parts_qe = self.__generate_population_part(phaseBit_mix_qe_before, phaseBit_mix_qe,
                                                 self.N - self.N // 2)
        parts = parts_ps + parts_qe
        for part in parts:
            plant = Plant(part.copy(), self.fobj)  # 创建植物对象
            self.pop.append(plant)

    def crossover(self, parent1, parent2):
        # 单点交叉
        point_row = np.random.randint(1, self.rows - 1)
        point_col = np.random.randint(1, self.cols - 1)

        child1_top_left = parent1[:point_row, :point_col]
        child1_top_right = parent2[:point_row, point_col:]
        child1_bottom_left = parent2[point_row:, :point_col]
        child1_bottom_right = parent1[point_row:, point_col:]

        child1 = np.vstack((np.hstack((child1_top_left, child1_top_right)),
                            np.hstack((child1_bottom_left, child1_bottom_right))))

        child2_top_left = parent2[:point_row, :point_col]
        child2_top_right = parent1[:point_row, point_col:]
        child2_bottom_left = parent1[point_row:, :point_col]
        child2_bottom_right = parent2[point_row:, point_col:]

        child2 = np.vstack((np.hstack((child2_top_left, child2_top_right)),
                            np.hstack((child2_bottom_left, child2_bottom_right))))

        return child1, child2

    def mutation(self, individual, mutation_rate=0.01):
        # 变异操作
        mutated_individual = individual.copy()
        for i in range(individual.shape[0]):
            for j in range(individual.shape[1]):
                if np.random.rand() < mutation_rate:
                    mutated_individual[i, j] = 1 - mutated_individual[i, j]
        return mutated_individual

    def run(self, phaseBit_mix_ps_before, phaseBit_mix_ps, phaseBit_mix_qe_before, phaseBit_mix_qe):

        # """初始化返回值"""
        self.best_individual = phaseBit_mix_ps
        self.best_fitness = self.fobj(self.best_individual)
        self.best_individual_history = []
        self.best_fitness_history = []

        # 普通初始化
        # self.initialize_population()

        # 初始化
        self.initialize_population_group(phaseBit_mix_ps_before, phaseBit_mix_ps,
                                         phaseBit_mix_qe_before, phaseBit_mix_qe)

        for it in range(self.max_iterations):
            costs = [plant.cost for plant in self.pop]  # 获取所有植物的成本
            best_cost = min(costs)  # 最佳成本

            newpop = []

            for i in range(self.N):
                ii = i + 1 if i < self.N - 1 else 0  # 下一个索引，循环到最后则回到第一个

                beta_1 = 1 + np.random.rand() / 2  # beta值，见论文中算法1第8行

                if self.pop[i].cost < beta_1 * self.pop[0].cost:
                    new_position1, new_position2 = self.crossover(self.pop[i].position, self.pop[ii].position)

                    # 添加变异操作
                    new_position1 = self.mutation(new_position1)
                    new_position2 = self.mutation(new_position2)

                    newsol1 = Plant(new_position1, self.fobj)  # 创建新植物对象
                    newsol2 = Plant(new_position2, self.fobj)  # 创建新植物对象

                    newpop.append(newsol1)
                    newpop.append(newsol2)
                else:
                    new_position = self.mutation(self.pop[0].position)
                    newsol = Plant(new_position, self.fobj)  # 创建新植物对象
                    newpop.append(newsol)

            # 合并种群
            self.pop.extend(newpop)

            # 排序种群
            self.pop.sort(key=lambda x: x.cost)

            # 竞争排斥（删除多余成员）
            if len(self.pop) > self.N:
                self.pop = self.pop[:self.N]

            # 存储迄今为止找到的最佳解
            best_sol = self.pop[0]
            self.best_costs[it] = best_sol.cost
            self.convergence_curve.append(best_sol.cost)

            # 记录最佳适应度曲线
            if best_sol.cost < self.best_fitness:
                self.best_individual = best_sol.position.copy()
                self.best_fitness = best_sol.cost
            self.best_fitness_history.append(self.best_fitness)
            self.best_individual_history.append(self.best_individual)

            # 打印当前最优解
            # print(f"Iteration {iteration + 1}: Best Fitness = {self.best_fitness}")
            logger.info("iteration=%d: self.best_fitness=%f" % (it, self.best_fitness))
            logger.info("---------------------------------------------------------------------------")

        destination_fitness = self.pop[0].cost  # 目标适应度
        destination_position = self.pop[0].position  # 目标位置

        # return destination_fitness, destination_position, self.convergence_curve
        return self.best_individual, self.best_fitness, self.best_fitness_history, self.best_individual_history



def test():
    # 设置参数
    dimensions = (4, 4)
    population_size = 30
    iterations = 100
    elite_ratio = 0.2
    # 创建一个示例 FitnessMultiBeamTrace 实例
    bit_num = 8  # 根据实际情况设置比特数
    beam_num = 2  # 根据实际情况设置波束数
    fitness_function = FitnessMultiBeamTrace(bit_num, beam_num)

    # 初始化主瓣位置
    phaseBit_target = np.random.randint(2, size=(4, 4))
    fitness_function.init_main_lope_pos(phaseBit_target)

    # 定义初始相位比特矩阵
    phaseBit_ps_before = np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]])
    phaseBit_ps = np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [1, 1, 1, 1]])
    phaseBit_qe_before = np.array([[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]])
    phaseBit_qe = np.array([[1, 1, 1, 1], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]])

    # 创建并运行二进制常春藤优化算法实例
    bioa = BinaryIvyOptimization(N=30,
                                 max_iterations=100,
                                 rows=3, cols=3,
                                 fobj=fitness_function.fitness)

    best_individual, best_fitness, best_fitness_history, best_individual_history \
        = bioa.run(phaseBit_ps_before, phaseBit_ps, phaseBit_qe_before, phaseBit_qe)

    print(f"\nBest Position:\n{best_individual.reshape(dimensions)}")
    print(f"Best Score: {best_fitness}")


if __name__=="__main__":
    test()
